package com.dmj.sqldsl.executor.transaction;


import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Stack;
import java.util.function.Supplier;
import lombok.AllArgsConstructor;

@AllArgsConstructor
public class TransactionManager {

  public TransactionStatus getTransaction(TransactionDefinition definition) {
    if (TransactionThreadContext.isEmpty()) {
      DefaultTransactionStatus status = new DefaultTransactionStatus(true);
      Transaction transaction = Transaction.create(definition, status);
      TransactionThreadContext.put(status, transaction);
      return status;
    }
    if (Propagation.ISOLATE == definition.getPropagation()) {
      DefaultTransactionStatus status = new DefaultTransactionStatus(true);
      Transaction transaction = Transaction.create(definition, status);
      TransactionThreadContext.put(status, transaction);
      return status;
    }
    DefaultTransactionStatus intoStatus = new DefaultTransactionStatus(false);
    TransactionThreadContext.add(intoStatus);
    return intoStatus;
  }

  public void commit(TransactionStatus status) {
    if (status.isCompleted()) {
      throw new TransactionException("Can not commit finished transaction!");
    }
    Transaction transaction = TransactionThreadContext.getTransaction(status)
        .orElseThrow(() -> new TransactionException("No transaction found!"));
    if (!status.isNewTransaction()) {
      transaction.remove(status);
      TransactionThreadContext.removeStatus(status);
      return;
    }
    try {
      transaction.commit();
    } catch (SQLException e) {
      throw new TransactionException(e);
    } finally {
      transaction.release();
      TransactionThreadContext.removeTransaction(transaction);
    }
  }

  public void rollback(TransactionStatus status) {
    if (status.isCompleted()) {
      throw new TransactionException("Can not rollback finished transaction!");
    }
    Transaction transaction = TransactionThreadContext.getTransaction(status)
        .orElseThrow(() -> new TransactionException("No transaction found!"));
    try {
      transaction.rollback();
    } catch (SQLException e) {
      throw new TransactionException(e);
    } finally {
      transaction.release();
      TransactionThreadContext.removeTransaction(transaction);
    }
  }

  protected static boolean hasTransaction() {
    return !TransactionThreadContext.isEmpty();
  }

  protected static Connection getTransactionConnection(Supplier<Connection> getConnectionFunction) {
    Transaction transaction = TransactionThreadContext.peek()
        .orElseThrow(() -> new TransactionException("No transaction found!"));
    return transaction.getTransactionConnection(getConnectionFunction);
  }

  private static class TransactionThreadContext {

    protected static final ThreadLocal<Map<TransactionStatus, Transaction>>
        transactionMapThreadLocal = ThreadLocal.withInitial(HashMap::new);
    protected static final ThreadLocal<Stack<Transaction>>
        stackThreadLocal = ThreadLocal.withInitial(Stack::new);

    public static void put(TransactionStatus status, Transaction transaction) {
      Map<TransactionStatus, Transaction> transactionMap = transactionMapThreadLocal.get();
      Stack<Transaction> stack = stackThreadLocal.get();
      transactionMap.put(status, transaction);
      stack.push(transaction);
    }

    public static void add(DefaultTransactionStatus status) {
      Map<TransactionStatus, Transaction> transactionMap = transactionMapThreadLocal.get();
      Stack<Transaction> stack = stackThreadLocal.get();
      Transaction transaction = stack.peek();
      transactionMap.put(status, transaction);
      transaction.getStatuses().add(status);
    }

    public static boolean isEmpty() {
      return transactionMapThreadLocal.get().isEmpty();
    }

    public static Optional<Transaction> getTransaction(TransactionStatus status) {
      Map<TransactionStatus, Transaction> transactionMap = transactionMapThreadLocal.get();
      return Optional.ofNullable(transactionMap.get(status));
    }

    public static void removeStatus(TransactionStatus status) {
      Map<TransactionStatus, Transaction> transactionMap = transactionMapThreadLocal.get();
      transactionMap.remove(status);
    }

    public static void removeTransaction(Transaction transaction) {
      Map<TransactionStatus, Transaction> transactionMap = transactionMapThreadLocal.get();
      transaction.getStatuses().forEach(transactionMap::remove);
      Stack<Transaction> stack = stackThreadLocal.get();
      stack.remove(transaction);
    }

    public static Optional<Transaction> peek() {
      Stack<Transaction> stack = stackThreadLocal.get();
      return Optional.ofNullable(stack.peek());
    }
  }
}
